from typing import Optional, List, Dict, Any

from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

class ModelScopeWordSegmentationLLM(LLM):
    task: str
    model_id: str
    def _call(
            self,
            prompt: str,
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,
    ) -> str:
        """Run the LLM on the given input.

        Override this method to implement the LLM logic.

        Args:
            prompt: The prompt to generate from.
            stop: Stop words to use when generating. Model output is cut off at the
                first occurrence of any of the stop substrings.
                If stop tokens are not supported consider raising NotImplementedError.
            run_manager: Callback manager for the run.
            **kwargs: Arbitrary additional keyword arguments. These are usually passed
                to the model provider API call.

        Returns:
            The model output as a string. Actual completions SHOULD NOT include the prompt.
        """
        # Use the ModelScope pipeline to process the prompt
        _pipeline = pipeline(task=self.task, model=self.model_id)
        resp = _pipeline(prompt)
        # Return the segmentation result
        return ','.join(map(str, resp['output']))

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        """Return a dictionary of identifying parameters."""
        return {
            # The model name allows users to specify custom token counting
            # rules in LLM monitoring applications (e.g., in LangSmith users
            # can provide per token pricing for their model and monitor
            # costs for the given LLM.)
            "model_name": "ModelScopeWordSegmentationLLM",
        }

    @property
    def _llm_type(self) -> str:
        """Get the type of language model used by this chat model. Used for logging purposes only."""
        return "echo_custom"

word_segmentation_llm = ModelScopeWordSegmentationLLM(
    task=Tasks.word_segmentation,
    model_id='damo/nlp_structbert_word-segmentation_chinese-base'
)

# 测试自定义的 LLM
input_text = '今天天气很好，我们去公园玩吧！'
output = word_segmentation_llm.invoke(input_text)
print(output)  # 输出分词结果
